# # 基于截断策略的机器阅读理解任务实现(滑动窗口)可以处理超过模型最大长度的文本
import collections
import numpy as np
import torch
from datasets import DatasetDict
from transformers import AutoTokenizer, AutoModelForQuestionAnswering, TrainingArguments, Trainer, DefaultDataCollator
from transformers import pipeline
from cmrc_eval import evaluate_cmrc


# 清理GPU缓存并设置使用GPU设备0
torch.cuda.empty_cache()
torch.cuda.set_device(0)

# # Step1 数据集加载
# 如果可以联网，直接使用load_dataset进行加载
# datasets = load_dataset("cmrc2018", cache_dir="data")
# 如果无法联网，则使用下面的方式加载数据集
datasets = DatasetDict.load_from_disk("./mrc_data")  # 从本地加载数据集
datasets
datasets["train"][0]  # 查看训练集第一条数据
print("-----------------------------1----------------------------------")

# # Step2 数据预处理
# 加载预训练的中文MacBERT分词器
tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")
tokenizer

# 选取前10个样本进行测试
sample_dataset = datasets["train"].select(range(10))

# 对样本进行分词处理，使用滑动窗口策略处理长文本
tokenized_examples = tokenizer(
	text = sample_dataset["question"],  # 问题文本
	text_pair = sample_dataset["context"],  # 上下文文本
	return_offsets_mapping = True,  # 返回字符偏移映射
	return_overflowing_tokens = True,  # 允许溢出token处理
	stride = 128,  # 滑动窗口步长
	max_length = 384,  # 最大序列长度
	truncation = "only_second",  # 只截断第二个文本(上下文)
	padding = "max_length"  # 填充到最大长度
)
tokenized_examples.keys()  # 查看返回的键值

# 查看样本映射关系和总数量
tokenized_examples["overflow_to_sample_mapping"], len(tokenized_examples["overflow_to_sample_mapping"])

# 解码并打印前3个样本
for sen in tokenizer.batch_decode(tokenized_examples["input_ids"][:3]):
	print(sen)

# 打印偏移映射信息
print(tokenized_examples["offset_mapping"][:3])

# 打印第一个样本的偏移映射及其长度
print(tokenized_examples["offset_mapping"][0], len(tokenized_examples["offset_mapping"][0]))

# 移除溢出到样本的映射关系
sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

# 处理每个分词后的样本，定位答案位置
for idx, _ in enumerate(sample_mapping):
	# 获取原始答案信息
	answer = sample_dataset["answers"][sample_mapping[idx]]
	start_char = answer["answer_start"][0]  # 答案起始字符位置
	end_char = start_char + len(answer["text"][0])  # 答案结束字符位置

	# 定位答案在token中的起始位置和结束位置
	# 一种策略，我们要拿到context的起始和结束，然后从左右两侧向答案逼近
	context_start = tokenized_examples.sequence_ids(idx).index(1)  # context起始位置
	context_end = tokenized_examples.sequence_ids(idx).index(None, context_start) - 1  # context结束位置

	offset = tokenized_examples.get("offset_mapping")[idx]
	example_ids = []

	# 判断答案是否在context中
	if offset[context_end][1] < start_char or offset[context_start][0] > end_char:
		start_token_pos = 0
		end_token_pos = 0
	else:
		# 从左侧向答案起始位置逼近
		token_id = context_start
		while token_id <= context_end and offset[token_id][0] < start_char:
			token_id += 1
		start_token_pos = token_id

		# 从右侧向答案结束位置逼近
		token_id = context_end
		while token_id >= context_start and offset[token_id][1] > end_char:
			token_id -= 1
		end_token_pos = token_id
		example_ids.append([sample_mapping[idx]])

	print(answer, start_char, end_char, context_start, context_end, start_token_pos, end_token_pos)
	print("token answer decode:",
		tokenizer.decode(tokenized_examples["input_ids"][idx][start_token_pos: end_token_pos + 1]))


# 数据预处理函数，用于整个数据集的批处理
def process_func(examples):
	# 对问题和上下文进行分词处理
	tokenized_examples = tokenizer(
		text = examples["question"],
		text_pair = examples["context"],
		return_offsets_mapping = True,
		return_overflowing_tokens = True,
		stride = 128,
		max_length = 384,
		truncation = "only_second",
		padding = "max_length"
	)

	# 获取样本映射关系
	sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
	start_positions = []  # 起始位置列表
	end_positions = []  # 结束位置列表
	example_ids = []  # 样本ID列表

	# 处理每个分词后的样本
	for idx, _ in enumerate(sample_mapping):
		# 获取答案信息
		answer = examples["answers"][sample_mapping[idx]]
		start_char = answer["answer_start"][0]
		end_char = start_char + len(answer["text"][0])

		# 获取context的起始和结束位置
		context_start = tokenized_examples.sequence_ids(idx).index(1)
		context_end = tokenized_examples.sequence_ids(idx).index(None, context_start) - 1
		offset = tokenized_examples.get("offset_mapping")[idx]

		# 判断答案是否在context中并定位答案token位置
		if offset[context_end][1] < start_char or offset[context_start][0] > end_char:
			start_token_pos = 0
			end_token_pos = 0
		else:
			# 定位起始位置
			token_id = context_start
			while token_id <= context_end and offset[token_id][0] < start_char:
				token_id += 1
			start_token_pos = token_id

			# 定位结束位置
			token_id = context_end
			while token_id >= context_start and offset[token_id][1] > end_char:
				token_id -= 1
			end_token_pos = token_id

		# 添加位置信息和样本ID
		start_positions.append(start_token_pos)
		end_positions.append(end_token_pos)
		example_ids.append(examples["id"][sample_mapping[idx]])

		# 更新偏移映射，只保留context部分
		tokenized_examples["offset_mapping"][idx] = [
			(o if tokenized_examples.sequence_ids(idx)[k] == 1 else None)
			for k, o in enumerate(tokenized_examples["offset_mapping"][idx])
		]

	# 添加处理后的字段
	tokenized_examples["example_ids"] = example_ids
	tokenized_examples["start_positions"] = start_positions
	tokenized_examples["end_positions"] = end_positions
	return tokenized_examples


# 对整个数据集应用预处理函数
tokenied_datasets = datasets.map(process_func, batched = True, remove_columns = datasets["train"].column_names)
tokenied_datasets

# 打印训练集第一个样本的偏移映射
print(tokenied_datasets["train"]["offset_mapping"][1])

# 打印前10个样本的ID
tokenied_datasets["train"]["example_ids"][:10]

# 建立example和feature的映射关系
example_to_feature = collections.defaultdict(list)
for idx, example_id in enumerate(tokenied_datasets["train"]["example_ids"][:10]):
	example_to_feature[example_id].append(idx)
example_to_feature
print("---------------------------------2------------------------------------------")


# # Step3 获取模型输出
# 处理模型预测结果的函数
def get_result(start_logits, end_logits, exmaples, features):
	predictions = {}  # 预测结果
	references = {}  # 参考答案

	# 建立example和feature的映射关系
	example_to_feature = collections.defaultdict(list)
	for idx, example_id in enumerate(features["example_ids"]):
		example_to_feature[example_id].append(idx)

	# 最优答案候选数量和最大答案长度
	n_best = 20
	max_answer_length = 30

	# 处理每个样本
	for example in exmaples:
		example_id = example["id"]
		context = example["context"]
		answers = []

		# 遍历该样本对应的所有feature
		for feature_idx in example_to_feature[example_id]:
			start_logit = start_logits[feature_idx]
			end_logit = end_logits[feature_idx]
			offset = features[feature_idx]["offset_mapping"]

			# 获取top n_best的起始和结束位置索引
			start_indexes = np.argsort(start_logit)[::-1][:n_best].tolist()
			end_indexes = np.argsort(end_logit)[::-1][:n_best].tolist()

			# 遍历所有可能的起始和结束位置组合
			for start_index in start_indexes:
				for end_index in end_indexes:
					# 过滤无效位置
					if offset[start_index] is None or offset[end_index] is None:
						continue
					if end_index < start_index or end_index - start_index + 1 > max_answer_length:
						continue

					# 添加候选答案
					answers.append({
						"text": context[offset[start_index][0]: offset[end_index][1]],
						"score": start_logit[start_index] + end_logit[end_index]
					})

		# 选择得分最高的答案
		if len(answers) > 0:
			best_answer = max(answers, key = lambda x: x["score"])
			predictions[example_id] = best_answer["text"]
		else:
			predictions[example_id] = ""

		references[example_id] = example["answers"]["text"]

	return predictions, references


print("---------------------------------3------------------------------------------")


# # Step4 评估函数
# 模型评估函数
def metirc(pred):
	# 获取起始和结束位置的logits
	start_logits, end_logits = pred[0]

	# 根据logits长度判断是验证集还是测试集
	if start_logits.shape[0] == len(tokenied_datasets["validation"]):
		p, r = get_result(start_logits, end_logits, datasets["validation"], tokenied_datasets["validation"])
	else:
		p, r = get_result(start_logits, end_logits, datasets["test"], tokenied_datasets["test"])

	# 使用CMRC评估函数进行评估
	return evaluate_cmrc(p, r)


print("---------------------------------4------------------------------------------")

# # Step5 加载模型
# 加载预训练的问答模型
model = AutoModelForQuestionAnswering.from_pretrained("hfl/chinese-macbert-base")


# 确保模型权重在内存中连续存储
def ensure_weights_contiguous(model):
	# 遍历模型所有命名参数
	for name, param in model.named_parameters():
		# 检查参数在内存中是否连续存储
		if not param.is_contiguous():
			print(f"Making {name} contiguous.")
			# 创建连续存储版本并替换原数据
			param.data = param.data.contiguous()


print("---------------------------------5------------------------------------------")
ensure_weights_contiguous(model)

# # Step6 配置TrainingArguments
# 设置训练参数
args = TrainingArguments(
	output_dir = "models_for_qa",  # 输出目录
	per_device_train_batch_size = 64,  # 训练时每个设备的批次大小
	per_device_eval_batch_size = 128,  # 评估时每个设备的批次大小
	gradient_accumulation_steps = 4,  # 梯度累积步数
	gradient_checkpointing = True,  # 启用梯度检查点以节省内存
	gradient_checkpointing_kwargs = {"use_reentrant": False},  # 推荐设置
	optim = "adafactor",  # 使用Adafactor优化器
	fp16 = True,  # 使用16位浮点数训练
	eval_strategy = "steps",  # 评估策略
	eval_steps = 200,  # 每200步评估一次
	save_strategy = "epoch",  # 保存策略：每个epoch保存一次
	logging_steps = 50,  # 日志记录步数
	num_train_epochs = 1  # 训练轮数
)
print("---------------------------------6------------------------------------------")

# # Step7 配置Trainer
# 配置训练器
trainer = Trainer(
	model = model,  # 模型
	args = args,  # 训练参数
	tokenizer = tokenizer,  # 分词器
	train_dataset = tokenied_datasets["train"],  # 训练数据集
	eval_dataset = tokenied_datasets["validation"],  # 验证数据集
	data_collator = DefaultDataCollator(),  # 数据整理器
	compute_metrics = metirc  # 评估函数
)
print("---------------------------------7------------------------------------------")

# # Step8 模型训练
# 开始训练
trainer.train()
print("---------------------------------8------------------------------------------")

# # Step9 模型预测
# 使用训练好的模型处理长文本
pipe = pipeline("question-answering", model = model, tokenizer = tokenizer, device = 0)

# 长文本示例 - 关于人工智能发展历史
long_context = """
人工智能（Artificial Intelligence，AI）是计算机科学的一个分支，旨在创建能够执行通常需要人类智能的任务的机器。人工智能的发展可以追溯到20世纪中期。

1950年，英国数学家阿兰·图灵提出了著名的图灵测试，用来判断机器是否具有智能。图灵认为，如果一台机器能够与人类进行对话而不被识别为机器，那么就可以说这台机器具有智能。

1956年，在美国达特茅斯学院举行的一次会议上，约翰·麦卡锡首次提出了"人工智能"这个术语，标志着人工智能作为一个独立学科的诞生。这次会议被认为是人工智能发展史上的里程碑事件。

20世纪60年代到70年代，人工智能研究主要集中在符号推理和专家系统上。研究人员试图通过编写规则来模拟人类的推理过程。1965年，约瑟夫·魏岑鲍姆开发了ELIZA程序，这是最早的聊天机器人之一，能够模拟心理治疗师的对话。

然而，到了20世纪70年代末和80年代初，人工智能遭遇了第一次"AI寒冬"。由于技术限制和过高的期望，研究资金大幅减少，人工智能的发展陷入停滞。

20世纪80年代，专家系统的兴起带来了人工智能的复兴。专家系统能够模拟人类专家的决策过程，在医疗诊断、地质勘探等领域取得了成功。日本的第五代计算机项目也推动了人工智能研究的发展。

21世纪初，随着计算能力的提升和大数据的出现，机器学习特别是深度学习技术开始兴起。2006年，杰弗里·辛顿提出了深度学习的概念，通过多层神经网络实现了更强大的模式识别能力。

2012年，亚历克斯·克里泽夫斯基等人开发的AlexNet在ImageNet图像识别竞赛中取得了突破性成绩，标志着深度学习在计算机视觉领域的成功应用。

2016年，谷歌DeepMind开发的AlphaGo程序击败了世界围棋冠军李世石，这是人工智能发展史上的又一重要里程碑，展示了深度学习和强化学习在复杂策略游戏中的强大能力。

如今，人工智能已经广泛应用于自然语言处理、计算机视觉、语音识别、推荐系统等各个领域，成为推动社会进步和经济发展的重要技术力量。
"""

# 提出一些需要在长文本中寻找答案的问题
questions = [
	"图灵测试是用来做什么的？",
	"谁首次提出了'人工智能'这个术语？",
	"ELIZA程序有什么特点？",
	"AlphaGo在哪一年击败了世界围棋冠军？",
	"什么是专家系统？"
]

# 对每个问题进行预测
for question in questions:
	result = pipe(question = question, context = long_context)
	print(f"问题: {question}")
	print(f"答案: {result['answer']}")
	print(f"置信度: {result['score']:.4f}")
	print("-" * 50)
print("---------------------------------9------------------------------------------")
